#!/bin/bash

# AIQuant 生产环境部署脚本
# 用途: 自动化生产环境部署，支持零停机更新
# 作者: AIQuant Team
# 版本: 1.0.0

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 配置变量
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
DEPLOY_ENV="production"
VERSION=${1:-latest}
HEALTH_CHECK_URL="https://api.aiquant.com/health"
BACKUP_DIR="/opt/backups/$(date +%Y%m%d_%H%M%S)"
MAX_HEALTH_CHECK_ATTEMPTS=60
HEALTH_CHECK_INTERVAL=10

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 错误处理
error_exit() {
    log_error "$1"
    exit 1
}

# 清理函数
cleanup() {
    log_info "执行清理操作..."
    # 清理临时文件
    rm -f /tmp/aiquant_deploy_*
}

# 注册清理函数
trap cleanup EXIT

# 显示帮助信息
show_help() {
    cat << EOF
AIQuant 生产环境部署脚本

用法: $0 [VERSION] [OPTIONS]

参数:
    VERSION     部署版本 (默认: latest)

选项:
    -h, --help              显示帮助信息
    -c, --check-only        仅执行预检查，不进行部署
    -b, --backup-only       仅创建备份
    -r, --rollback VERSION  回滚到指定版本
    -f, --force             强制部署，跳过确认
    --skip-backup           跳过备份步骤
    --skip-tests            跳过部署后测试

示例:
    $0 v1.2.3                    # 部署v1.2.3版本
    $0 --check-only              # 仅执行预检查
    $0 --rollback v1.2.2         # 回滚到v1.2.2
    $0 v1.2.3 --skip-backup      # 部署但跳过备份

EOF
}

# 解析命令行参数
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -c|--check-only)
                CHECK_ONLY=true
                shift
                ;;
            -b|--backup-only)
                BACKUP_ONLY=true
                shift
                ;;
            -r|--rollback)
                ROLLBACK_VERSION="$2"
                shift 2
                ;;
            -f|--force)
                FORCE_DEPLOY=true
                shift
                ;;
            --skip-backup)
                SKIP_BACKUP=true
                shift
                ;;
            --skip-tests)
                SKIP_TESTS=true
                shift
                ;;
            -*)
                error_exit "未知选项: $1"
                ;;
            *)
                if [[ -z "$VERSION" || "$VERSION" == "latest" ]]; then
                    VERSION="$1"
                fi
                shift
                ;;
        esac
    done
}

# 预检查
pre_deploy_check() {
    log_info "执行部署前检查..."
    
    # 检查必需的命令
    local required_commands=("docker" "docker-compose" "curl" "jq" "pg_dump" "redis-cli")
    for cmd in "${required_commands[@]}"; do
        if ! command -v "$cmd" &> /dev/null; then
            error_exit "缺少必需命令: $cmd"
        fi
    done
    
    # 检查Docker服务
    if ! docker info &> /dev/null; then
        error_exit "Docker服务未运行"
    fi
    
    # 检查磁盘空间
    local available_space=$(df / | awk 'NR==2{print $4}')
    local required_space=10485760  # 10GB in KB
    if [[ $available_space -lt $required_space ]]; then
        error_exit "磁盘空间不足，需要至少10GB可用空间"
    fi
    
    # 检查内存
    local available_memory=$(free | awk 'NR==2{print $7}')
    local required_memory=4194304  # 4GB in KB
    if [[ $available_memory -lt $required_memory ]]; then
        log_warning "可用内存不足4GB，部署可能会失败"
    fi
    
    # 检查网络连接
    if ! curl -s --connect-timeout 5 https://registry.hub.docker.com &> /dev/null; then
        log_warning "无法连接到Docker Hub，可能影响镜像拉取"
    fi
    
    # 检查配置文件
    if [[ ! -f "$PROJECT_ROOT/.env.production" ]]; then
        error_exit "生产环境配置文件不存在: .env.production"
    fi
    
    # 检查Docker Compose文件
    if [[ ! -f "$PROJECT_ROOT/docker-compose.prod.yml" ]]; then
        error_exit "生产环境Docker Compose文件不存在"
    fi
    
    # 验证Docker Compose配置
    cd "$PROJECT_ROOT"
    if ! docker-compose -f docker-compose.prod.yml config &> /dev/null; then
        error_exit "Docker Compose配置文件有误"
    fi
    
    log_success "预检查完成"
}

# 创建备份
create_backup() {
    if [[ "$SKIP_BACKUP" == "true" ]]; then
        log_info "跳过备份步骤"
        return 0
    fi
    
    log_info "创建系统备份到: $BACKUP_DIR"
    mkdir -p "$BACKUP_DIR"
    
    # 备份数据库
    log_info "备份PostgreSQL数据库..."
    if docker-compose -f docker-compose.prod.yml exec -T postgres pg_isready -U aiquant; then
        docker-compose -f docker-compose.prod.yml exec -T postgres \
            pg_dump -U aiquant -d aiquant --verbose --format=custom \
            > "$BACKUP_DIR/database.dump" || error_exit "数据库备份失败"
        
        # 压缩数据库备份
        gzip "$BACKUP_DIR/database.dump"
        log_success "数据库备份完成"
    else
        log_warning "数据库不可用，跳过数据库备份"
    fi
    
    # 备份Redis数据
    log_info "备份Redis数据..."
    if docker-compose -f docker-compose.prod.yml exec -T redis redis-cli ping | grep -q PONG; then
        docker-compose -f docker-compose.prod.yml exec -T redis \
            redis-cli --rdb /tmp/redis_backup.rdb
        docker-compose -f docker-compose.prod.yml exec -T redis \
            cat /tmp/redis_backup.rdb > "$BACKUP_DIR/redis.rdb" || log_warning "Redis备份失败"
        log_success "Redis备份完成"
    else
        log_warning "Redis不可用，跳过Redis备份"
    fi
    
    # 备份配置文件
    log_info "备份配置文件..."
    tar -czf "$BACKUP_DIR/config.tar.gz" \
        -C "$PROJECT_ROOT" \
        .env.production \
        docker-compose.prod.yml \
        nginx/ \
        ssl/ \
        monitoring/ 2>/dev/null || log_warning "部分配置文件备份失败"
    
    # 备份应用数据
    log_info "备份应用数据..."
    if [[ -d "$PROJECT_ROOT/data" ]]; then
        tar -czf "$BACKUP_DIR/appdata.tar.gz" \
            -C "$PROJECT_ROOT" data/ || log_warning "应用数据备份失败"
    fi
    
    # 创建备份清单
    cat > "$BACKUP_DIR/backup_info.txt" << EOF
备份时间: $(date)
部署版本: $VERSION
备份目录: $BACKUP_DIR
系统信息: $(uname -a)
Docker版本: $(docker --version)
EOF
    
    # 上传备份到云存储（如果配置了）
    if [[ -n "$AWS_S3_BACKUP_BUCKET" ]]; then
        log_info "上传备份到S3..."
        aws s3 sync "$BACKUP_DIR/" "s3://$AWS_S3_BACKUP_BUCKET/backups/$(basename $BACKUP_DIR)/" || log_warning "S3上传失败"
    fi
    
    log_success "备份创建完成: $BACKUP_DIR"
}

# 拉取新镜像
pull_images() {
    log_info "拉取Docker镜像 - 版本: $VERSION"
    
    cd "$PROJECT_ROOT"
    
    # 设置镜像标签
    export IMAGE_TAG="$VERSION"
    
    # 拉取镜像
    local images=("aiquant:$VERSION" "aiquant-web:$VERSION" "aiquant-worker:$VERSION")
    
    for image in "${images[@]}"; do
        log_info "拉取镜像: $image"
        if ! docker pull "$image"; then
            error_exit "拉取镜像失败: $image"
        fi
    done
    
    # 验证镜像
    for image in "${images[@]}"; do
        if ! docker inspect "$image" &> /dev/null; then
            error_exit "镜像验证失败: $image"
        fi
    done
    
    log_success "镜像拉取完成"
}

# 执行滚动更新
rolling_update() {
    log_info "执行滚动更新..."
    
    cd "$PROJECT_ROOT"
    
    # 设置环境变量
    export IMAGE_TAG="$VERSION"
    export COMPOSE_FILE="docker-compose.prod.yml"
    
    # 更新服务（一次更新一个实例）
    local services=("api" "web" "worker")
    
    for service in "${services[@]}"; do
        log_info "更新服务: $service"
        
        # 获取当前副本数
        local current_replicas=$(docker-compose ps -q "$service" | wc -l)
        
        if [[ $current_replicas -gt 1 ]]; then
            # 多副本服务，逐个更新
            for ((i=1; i<=current_replicas; i++)); do
                log_info "更新 $service 实例 $i/$current_replicas"
                
                # 停止一个实例
                local container_id=$(docker-compose ps -q "$service" | head -1)
                if [[ -n "$container_id" ]]; then
                    docker stop "$container_id"
                    docker rm "$container_id"
                fi
                
                # 启动新实例
                docker-compose up -d --scale "$service=$current_replicas" "$service"
                
                # 等待实例就绪
                sleep 30
                
                # 健康检查
                if ! check_service_health "$service"; then
                    error_exit "$service 服务更新失败"
                fi
            done
        else
            # 单副本服务，直接更新
            docker-compose up -d "$service"
            sleep 30
            if ! check_service_health "$service"; then
                error_exit "$service 服务更新失败"
            fi
        fi
        
        log_success "$service 服务更新完成"
    done
    
    log_success "滚动更新完成"
}

# 检查服务健康状态
check_service_health() {
    local service="$1"
    local max_attempts=30
    local attempt=1
    
    log_info "检查 $service 服务健康状态..."
    
    while [[ $attempt -le $max_attempts ]]; do
        if docker-compose ps "$service" | grep -q "Up"; then
            # 检查容器健康状态
            local container_id=$(docker-compose ps -q "$service" | head -1)
            if [[ -n "$container_id" ]]; then
                local health_status=$(docker inspect --format='{{.State.Health.Status}}' "$container_id" 2>/dev/null || echo "none")
                if [[ "$health_status" == "healthy" || "$health_status" == "none" ]]; then
                    log_success "$service 服务健康"
                    return 0
                fi
            fi
        fi
        
        log_info "等待 $service 服务就绪... ($attempt/$max_attempts)"
        sleep 10
        ((attempt++))
    done
    
    log_error "$service 服务健康检查失败"
    return 1
}

# 系统健康检查
system_health_check() {
    log_info "执行系统健康检查..."
    
    local attempt=1
    while [[ $attempt -le $MAX_HEALTH_CHECK_ATTEMPTS ]]; do
        log_info "健康检查尝试 $attempt/$MAX_HEALTH_CHECK_ATTEMPTS"
        
        # 检查API健康状态
        if curl -f -s --connect-timeout 5 --max-time 10 "$HEALTH_CHECK_URL" > /dev/null; then
            log_success "API健康检查通过"
            
            # 检查数据库连接
            if docker-compose -f docker-compose.prod.yml exec -T postgres pg_isready -U aiquant; then
                log_success "数据库连接正常"
            else
                log_warning "数据库连接异常"
            fi
            
            # 检查Redis连接
            if docker-compose -f docker-compose.prod.yml exec -T redis redis-cli ping | grep -q PONG; then
                log_success "Redis连接正常"
            else
                log_warning "Redis连接异常"
            fi
            
            return 0
        fi
        
        log_info "等待系统就绪..."
        sleep $HEALTH_CHECK_INTERVAL
        ((attempt++))
    done
    
    error_exit "系统健康检查失败"
}

# 数据库迁移
run_database_migration() {
    log_info "执行数据库迁移..."
    
    cd "$PROJECT_ROOT"
    
    # 等待数据库就绪
    local attempt=1
    while [[ $attempt -le 30 ]]; do
        if docker-compose -f docker-compose.prod.yml exec -T postgres pg_isready -U aiquant; then
            break
        fi
        log_info "等待数据库就绪... ($attempt/30)"
        sleep 5
        ((attempt++))
    done
    
    # 执行迁移
    if docker-compose -f docker-compose.prod.yml exec -T api python scripts/migrate_db.py; then
        log_success "数据库迁移完成"
    else
        log_warning "数据库迁移失败，请手动检查"
    fi
}

# 缓存预热
warm_cache() {
    log_info "执行缓存预热..."
    
    cd "$PROJECT_ROOT"
    
    if docker-compose -f docker-compose.prod.yml exec -T api python scripts/warm_cache.py; then
        log_success "缓存预热完成"
    else
        log_warning "缓存预热失败"
    fi
}

# 部署后测试
post_deploy_tests() {
    if [[ "$SKIP_TESTS" == "true" ]]; then
        log_info "跳过部署后测试"
        return 0
    fi
    
    log_info "执行部署后测试..."
    
    cd "$PROJECT_ROOT"
    
    # 基本功能测试
    log_info "执行基本功能测试..."
    if python scripts/test_basic_functions.py; then
        log_success "基本功能测试通过"
    else
        log_warning "基本功能测试失败"
    fi
    
    # API端点测试
    log_info "测试关键API端点..."
    local endpoints=(
        "/health"
        "/api/v1/health"
        "/api/v1/market/status"
        "/api/v1/user/profile"
    )
    
    for endpoint in "${endpoints[@]}"; do
        if curl -f -s "${HEALTH_CHECK_URL%/health}$endpoint" > /dev/null; then
            log_success "API端点测试通过: $endpoint"
        else
            log_warning "API端点测试失败: $endpoint"
        fi
    done
}

# 清理旧镜像
cleanup_old_images() {
    log_info "清理旧Docker镜像..."
    
    # 清理未使用的镜像
    docker image prune -f
    
    # 清理旧版本镜像（保留最近3个版本）
    local images=("aiquant" "aiquant-web" "aiquant-worker")
    for image in "${images[@]}"; do
        local old_images=$(docker images "$image" --format "table {{.Repository}}:{{.Tag}}" | tail -n +2 | tail -n +4)
        if [[ -n "$old_images" ]]; then
            echo "$old_images" | xargs -r docker rmi || log_warning "清理旧镜像失败: $image"
        fi
    done
    
    log_success "镜像清理完成"
}

# 发送通知
send_notification() {
    local status="$1"
    local message="$2"
    
    # Slack通知
    if [[ -n "$SLACK_WEBHOOK_URL" ]]; then
        local color="good"
        if [[ "$status" == "error" ]]; then
            color="danger"
        elif [[ "$status" == "warning" ]]; then
            color="warning"
        fi
        
        curl -X POST "$SLACK_WEBHOOK_URL" \
            -H 'Content-type: application/json' \
            --data "{
                \"attachments\": [{
                    \"color\": \"$color\",
                    \"title\": \"AIQuant生产环境部署\",
                    \"text\": \"$message\",
                    \"fields\": [
                        {\"title\": \"版本\", \"value\": \"$VERSION\", \"short\": true},
                        {\"title\": \"环境\", \"value\": \"$DEPLOY_ENV\", \"short\": true},
                        {\"title\": \"时间\", \"value\": \"$(date)\", \"short\": false}
                    ]
                }]
            }" || log_warning "Slack通知发送失败"
    fi
    
    # 邮件通知
    if [[ -n "$EMAIL_RECIPIENTS" ]]; then
        echo "$message" | mail -s "AIQuant部署通知 - $status" "$EMAIL_RECIPIENTS" || log_warning "邮件通知发送失败"
    fi
}

# 回滚功能
rollback() {
    local rollback_version="$1"
    
    log_info "开始回滚到版本: $rollback_version"
    
    # 检查回滚版本是否存在
    if ! docker inspect "aiquant:$rollback_version" &> /dev/null; then
        error_exit "回滚版本不存在: $rollback_version"
    fi
    
    # 设置回滚版本
    export IMAGE_TAG="$rollback_version"
    
    # 执行回滚
    cd "$PROJECT_ROOT"
    docker-compose -f docker-compose.prod.yml up -d
    
    # 健康检查
    if system_health_check; then
        log_success "回滚完成: $rollback_version"
        send_notification "success" "系统已成功回滚到版本 $rollback_version"
    else
        error_exit "回滚失败"
    fi
}

# 主函数
main() {
    log_info "开始AIQuant生产环境部署"
    log_info "版本: $VERSION"
    log_info "环境: $DEPLOY_ENV"
    log_info "时间: $(date)"
    
    # 解析命令行参数
    parse_args "$@"
    
    # 处理特殊操作
    if [[ "$BACKUP_ONLY" == "true" ]]; then
        create_backup
        exit 0
    fi
    
    if [[ -n "$ROLLBACK_VERSION" ]]; then
        rollback "$ROLLBACK_VERSION"
        exit 0
    fi
    
    # 预检查
    pre_deploy_check
    
    if [[ "$CHECK_ONLY" == "true" ]]; then
        log_success "预检查完成，系统准备就绪"
        exit 0
    fi
    
    # 确认部署
    if [[ "$FORCE_DEPLOY" != "true" ]]; then
        echo
        log_warning "即将部署到生产环境！"
        echo "版本: $VERSION"
        echo "环境: $DEPLOY_ENV"
        echo "备份目录: $BACKUP_DIR"
        echo
        read -p "确认继续部署？(y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "部署已取消"
            exit 0
        fi
    fi
    
    # 记录部署开始时间
    local deploy_start_time=$(date +%s)
    
    # 执行部署步骤
    create_backup
    pull_images
    rolling_update
    run_database_migration
    system_health_check
    warm_cache
    post_deploy_tests
    cleanup_old_images
    
    # 计算部署时间
    local deploy_end_time=$(date +%s)
    local deploy_duration=$((deploy_end_time - deploy_start_time))
    
    log_success "生产环境部署完成！"
    log_info "部署版本: $VERSION"
    log_info "部署时间: ${deploy_duration}秒"
    log_info "备份位置: $BACKUP_DIR"
    
    # 发送成功通知
    send_notification "success" "AIQuant生产环境部署成功完成！版本: $VERSION，耗时: ${deploy_duration}秒"
}

# 错误处理
set -E
trap 'log_error "部署过程中发生错误，行号: $LINENO"; send_notification "error" "AIQuant生产环境部署失败！版本: $VERSION"; exit 1' ERR

# 执行主函数
main "$@"